package com.nbufe.utils;

import org.apache.commons.math3.analysis.polynomials.PolynomialFunction;
import org.apache.commons.math3.fitting.PolynomialCurveFitter;
import org.apache.commons.math3.fitting.WeightedObservedPoints;
import org.python.core.Py;
import org.python.core.PyObject;
import org.python.util.PythonInterpreter;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 自动回归预测模型
 *
 * @author luzhiheng
 * @date 2024-01-18
 */
public class AutoRegressionModelUtil {

    /**
     * 使用多项式回归和季节性因素进行自回归预测
     *
     * @param data         输入数据
     * @param m            预测未来m个数据点
     * @param order        多项式回归的阶数
     * @param seasonalData 季节性因素数据
     * @return 预测结果数组
     */
    public static double[] autoRegressionWithSeasonalityPrediction(double[] data, int m, int order, double[] seasonalData) {
        int n = data.length;

        // 检查输入数据合法性
        if (n < order + 1 || n != seasonalData.length) {
            throw new IllegalArgumentException("Input data and seasonal data sizes must be consistent.");
        }

        // 构建观测点
        WeightedObservedPoints obs = new WeightedObservedPoints();
        for (int i = 0; i < n; i++) {
            obs.add(i, data[i]);
        }

        try {
            // 拟合多项式曲线
            PolynomialCurveFitter fitter = PolynomialCurveFitter.create(order);
            double[] coefficients = fitter.fit(obs.toList());

            // 预测未来m个数据点
            double[] forecastValues = new double[m];
            for (int i = 0; i < m; i++) {
                double forecast = evaluatePolynomial(coefficients, n + i) + seasonalData[(n + i) % seasonalData.length];
                forecast = FormatUtil.saveTwoDecimalPlace(forecast);
                forecastValues[i] = forecast;
            }

            return forecastValues;
        } catch (Exception e) {
            // 拟合过程中出现异常
            throw new RuntimeException("Error during curve fitting.", e);
        }
    }

    /**
     * 计算多项式回归函数值
     *
     * @param coefficients 多项式系数
     * @param x            自变量
     * @return 函数值
     */
    private static double evaluatePolynomial(double[] coefficients, double x) {
        PolynomialFunction polynomial = new PolynomialFunction(coefficients);
        return polynomial.value(x);
    }
}